from collections import defaultdict
from typing import Callable, Dict, List

from chernc.treediff.gumtree.matchers.heuristic.gt.priority_tree_queue import PriorityTreeQueue
from chernc.treediff.gumtree.tree.tree import TreeLike


class DefaultPriorityTreeQueue(PriorityTreeQueue):
    def __init__(self, root: TreeLike, minimum_priority: int, priority_calculator: Callable[[TreeLike], int]) -> None:
        self.trees: Dict[int, List[TreeLike]] = defaultdict(list)  # 使用 defaultdict 来模拟 Int2ObjectSortedMap
        self.minimum_priority: int = minimum_priority
        self.priority_calculator: Callable[[TreeLike], int] = priority_calculator
        self.add(root)

    def pop_open(self) -> List[TreeLike]:
        pop = self.pop()
        for t in pop:
            self.open(t)
        return pop

    def set_priority_calculator(self, calculator: Callable[[TreeLike], int]) -> None:
        self.priority_calculator = calculator

    def pop(self) -> List[TreeLike]:
        current_priority = self.current_priority()
        return self.trees.pop(current_priority, [])

    def open(self, tree: TreeLike) -> None:
        for c in tree.get_children():
            self.add(c)

    def current_priority(self) -> int:
        if not self.trees:
            return -1  # 默认返回一个表示空队列的值
        return max(self.trees.keys())  # 获取字典中的最大键

    def set_minimum_priority(self, priority: int) -> None:
        self.minimum_priority = priority

    def get_minimum_priority(self) -> int:
        return self.minimum_priority

    def is_empty(self) -> bool:
        return not bool(self.trees)

    def clear(self) -> None:
        self.trees.clear()

    def add(self, tree: TreeLike) -> None:
        priority = self.priority_calculator(tree)
        if priority < self.minimum_priority:
            return

        if priority not in self.trees:
            self.trees[priority] = []
        self.trees[priority].append(tree)
